package mc.support.zmq.route;

import mc.support.proto.ZmqRegisterReq;
import mc.support.proto.ZmqRegisterRsp;
import mc.support.task.common.CronTaskAdaptor;
import mc.support.zmq.ProtoConstants;
import mc.support.zmq.err.NoAvailableServerException;
import org.slf4j.Logger;
import org.zeromq.ZMQ;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Created with IntelliJ IDEA. <br/>
 * Author: Francis Yuen    <br/>
 * Date: 2014-10-24  <br/>
 */
public class ZmqServerRouterImpl extends CronTaskAdaptor implements ZmqServerRouter, ZmqServerRegister {

    private static final Logger log = org.slf4j.LoggerFactory.getLogger( ZmqServerRouterImpl.class );

    private static final int ALLOWED_TIMEOUT = 3;
    public static final int MAGIC_CODE = ProtoConstants.MAGIC_CODE;

    private final int port;
    private final int ioThreads;
    private final AtomicBoolean started = new AtomicBoolean( false );
    private final AtomicBoolean keepOn = new AtomicBoolean( true );
    private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
    private final Lock rLock = rwLock.readLock();
    private final Lock wLock = rwLock.writeLock();
    private final Random r = new Random();
    private final Map<String, List<String>> servers = new HashMap<>();
    private final ZmqRegisterRsp dummyRsp = ZmqRegisterRsp
            .newBuilder().setMagicCode( MAGIC_CODE ).setSuccess( false ).build();

    private final Map<String, Integer> timeoutStatis = new HashMap<>();

    private final TimeoutHandler defaultTimeoutHandler = new TimeoutHandler() {
        @Override
        public void handle( ZmqConnection conn, long timeoutMs ) {
            synchronized ( timeoutStatis ) {
                String address = conn.getServerAddress();
                Integer count = timeoutStatis.get( address );
                if ( null == count ) count = 0;
                timeoutStatis.put( address, ++count );
                if ( count >= ALLOWED_TIMEOUT ) {
                    unregisterBizServer( conn.getBizType(), address );
                    timeoutStatis.remove( address );
                }
            }
        }
    };

    private ZMQ.Context context;

    public ZmqServerRouterImpl( int port, int ioThreads ) {
        this.port = port;
        this.ioThreads = ioThreads;
    }

    @Override
    public boolean isStarted() {
        return started.get();
    }

    @Override
    public void start() {
        context = ZMQ.context( ioThreads );
        ZMQ.Socket broker = context.socket( ZMQ.REP );
        broker.bind( "tcp://*:" + port );

        while ( keepOn.get() ) {
            try {
                byte[] bs = broker.recv();
                ZmqRegisterReq req = ZmqRegisterReq.parseFrom( bs );
                if ( MAGIC_CODE == req.getMagicCode() ) {
                    ZmqRegisterRsp rsp;
                    switch ( req.getOps() ) {
                        case REGISTER:
                            rsp = registerBizServer( req.getBizType(), req.getAddress() );
                            break;
                        case UNREGISTER:
                            rsp = unregisterBizServer( req.getBizType(), req.getAddress() );
                            break;
                        default:
                            rsp = dummyRsp;
                    }
                    broker.send( rsp.toByteArray() );
                } else {
                    log.warn( "magic code not match: ", req.getMagicCode() );
                }
            } catch ( Exception err ) {
                log.error( "error in zmq server router: ", err );
            }
        }

    }

    @Override
    public boolean isStopped() {
        if ( !isStarted() ) throw new IllegalStateException( "ZMQ server router not started!" );
        return !keepOn.get();
    }

    @Override
    public void stop() {
        if ( !isStarted() ) throw new IllegalStateException( "ZMQ server router not started!" );
        keepOn.set( false );
    }

    @Override
    public boolean hasAvailable( String bizType ) {
        try {
            rLock.lock();
            return !servers.isEmpty();
        } finally {
            rLock.unlock();
        }
    }

    @Override
    public ZmqConnection acquireOne( String bizType ) {
        return acquireOne( bizType, ZMQ.REQ );
    }

    @Override
    public ZmqConnection acquireOne( String bizType, int zmqType ) {
        String address = null;
        try {
            rLock.lock();
            List<String> addresses = servers.get( bizType );
            final int len = null != addresses ? addresses.size() : 0;
            if ( len > 0 ) {
                address = addresses.get( r.nextInt( len ) );
            }
        } finally {
            rLock.unlock();
        }
        if ( null != address ) {
            ZMQ.Socket socket = context.socket( zmqType );
            socket.connect( address );
            return new ZmqConnection( socket, bizType, address )
                    .setTimeoutHandler( defaultTimeoutHandler );
        }
        throw new NoAvailableServerException( "no available server for bizType: " + bizType );
    }

    public ZmqRegisterRsp registerBizServer( String bizType, String address ) {
        String msg = "";
        boolean success = false, registered;
        if ( registered = isRegistered( bizType, address ) ) {
            msg = address + " is already registered!";
            success = true;
        }
        if ( !registered ) {
            try {
                wLock.lock();
                List<String> addresses = servers.get( bizType );
                if ( null == addresses ) {
                    addresses = new ArrayList<>();
                    servers.put( bizType, addresses );
                }
                success = addresses.add( address );
                log.info( "server {} registered!", address );
                msg = address + " register succeed!";
            } finally {
                wLock.unlock();
            }
        }
        return ZmqRegisterRsp.newBuilder()
                .setMagicCode( MAGIC_CODE )
                .setSuccess( success )
                .setMsg( msg ).build();
    }

    private boolean isRegistered( String bizType, String address ) {
        try {
            rLock.lock();
            List<String> addresses = servers.get( bizType );
            return null != addresses && addresses.contains( address );
        } finally {
            rLock.unlock();
        }
    }

    public ZmqRegisterRsp unregisterBizServer( String bizType, String address ) {
        try {
            wLock.lock();
            String msg;
            boolean success = false;
            List<String> addresses = servers.get( bizType );
            if ( null != addresses ) {
                success = addresses.remove( address );
                msg = address + " unregister status: " + success;
                if ( 0 == addresses.size() ) {
                    servers.remove( bizType );
                    msg += ", no available server for bizType: " + bizType + " after unregister!";
                }
            } else {
                msg = "there is no available server for bizType: " + bizType;
            }
            return ZmqRegisterRsp.newBuilder()
                    .setMagicCode( MAGIC_CODE )
                    .setSuccess( success )
                    .setMsg( msg )
                    .build();
        } finally {
            wLock.unlock();
        }
    }

    @Override
    public boolean register( String bizType, String address ) {
        return registerBizServer( bizType, address ).getSuccess();
    }

    @Override
    public boolean unregister( String bizType, String address ) {
        return registerBizServer( bizType, address ).getSuccess();
    }

    @Override
    public void run() {
        start();
    }

    @Override
    public boolean isLaunchNow() {
        return true;
    }

}
